package com.knowledge.business.util;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.knowledge.business.config.SiliconFlowConfig;
import com.knowledge.business.model.EmbeddingRequest;
import com.knowledge.business.model.EmbeddingResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * 硅基流动向量模型工具类
 *
 * @author knowledge
 */
@Component
public class SiliconFlowEmbeddingUtil {

    private static final Logger log = LoggerFactory.getLogger(SiliconFlowEmbeddingUtil.class);

    @Autowired
    private SiliconFlowConfig siliconFlowConfig;

    private final HttpClient httpClient;
    private final ObjectMapper objectMapper;

    public SiliconFlowEmbeddingUtil() {
        this.httpClient = HttpClient.newBuilder()
                .connectTimeout(Duration.ofSeconds(30))
                .build();
        this.objectMapper = new ObjectMapper();
    }

    /**
     * 调用硅基流动向量模型生成向量
     *
     * @param text 输入文本
     * @return 向量数据
     * @throws IOException
     * @throws InterruptedException
     */
    public List<Float> generateEmbedding(String text) throws IOException, InterruptedException {
        int maxRetries = 3;
        int retryCount = 0;
        while (retryCount < maxRetries) {
            try {
                return attemptGenerateEmbedding(text);
            } catch (IOException e) {
                retryCount++;
                if (retryCount >= maxRetries) {
                    log.error("调用硅基流动向量模型失败，达到最大重试次数: {}", maxRetries, e);
                    throw e;
                }

                // 指数退避策略
                long backoffTime = (long) Math.pow(2, retryCount) * 1000;
                log.warn("调用硅基流动向量模型失败，{}ms后进行第{}次重试", backoffTime, retryCount, e);
                Thread.sleep(backoffTime);
            }
        }

        throw new RuntimeException("调用硅基流动向量模型失败，达到最大重试次数");
    }

    private List<Float> attemptGenerateEmbedding(String text) throws IOException, InterruptedException {
        // 构造请求
        EmbeddingRequest request = new EmbeddingRequest(
                siliconFlowConfig.getModel(),
                Collections.singletonList(text),
                "float"
        );

        String requestBody = objectMapper.writeValueAsString(request);

        // 构造HTTP请求
        HttpRequest httpRequest = HttpRequest.newBuilder()
                .uri(URI.create(siliconFlowConfig.getBaseUrl() + "/embeddings"))
                .header("Authorization", "Bearer " + siliconFlowConfig.getApiKey())
                .header("Content-Type", "application/json")
                .timeout(Duration.ofSeconds(30))  // 设置请求超时时间
                .POST(HttpRequest.BodyPublishers.ofString(requestBody))
                .build();

        // 发送请求
        HttpResponse<String> response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());

        if (response.statusCode() != 200) {
            log.error("调用硅基流动向量模型失败，状态码: {}，响应: {}", response.statusCode(), response.body());
            throw new RuntimeException("调用硅基流动向量模型失败，状态码: " + response.statusCode());
        }

        // 解析响应
        EmbeddingResponse embeddingResponse = objectMapper.readValue(response.body(), EmbeddingResponse.class);

        if (embeddingResponse.getData() == null || embeddingResponse.getData().isEmpty()) {
            throw new RuntimeException("向量模型返回数据为空");
        }

        return embeddingResponse.getData().get(0).getEmbedding();
    }

    /**
     * 将float列表转换为byte数组
     *
     * @param floatList float列表
     * @return byte数组
     */
    public byte[] floatListToByteArray(List<Float> floatList) {
        byte[] byteArray = new byte[floatList.size() * 4];
        for (int i = 0; i < floatList.size(); i++) {
            int bits = Float.floatToIntBits(floatList.get(i));
            byteArray[i * 4] = (byte) (bits & 0xff);
            byteArray[i * 4 + 1] = (byte) ((bits >> 8) & 0xff);
            byteArray[i * 4 + 2] = (byte) ((bits >> 16) & 0xff);
            byteArray[i * 4 + 3] = (byte) ((bits >> 24) & 0xff);
        }
        return byteArray;
    }

    /**
     * 将byte数组转换为Float列表
     *
     * @param bytes byte数组
     * @return Float列表
     */
    public List<Float> byteArrayToFloatList(byte[] bytes) {
        if (bytes == null || bytes.length == 0) {
            return new ArrayList<>();
        }

        List<Float> floatList = new ArrayList<>();
        for (int i = 0; i < bytes.length; i += 4) {
            int asInt = ((bytes[i] & 0xFF) << 24) |
                    ((bytes[i + 1] & 0xFF) << 16) |
                    ((bytes[i + 2] & 0xFF) << 8) |
                    (bytes[i + 3] & 0xFF);
            floatList.add(Float.intBitsToFloat(asInt));
        }
        return floatList;
    }
}
